# --------------------------------------------------------
# SiamMask
# Licensed under The MIT License
# Written by Qiang Wang (wangqiang2015 at ia.ac.cn)
# --------------------------------------------------------
import glob

#################################
import os #mod
import sys
if os.name == 'nt':
	path = os.getcwd()                #modified for windows version
  #print(path)                      
	sys.path.append(path)
	path, _ = os.path.split(path)
	path, _ = os.path.split(path)
	#print(path)
	sys.path.append(path)
#############################
from tools.test import *  

parser = argparse.ArgumentParser(description='PyTorch Tracking Demo')

parser.add_argument('--resume', default='', type=str, required=True,
                    metavar='PATH',help='path to latest checkpoint (default: none)')
parser.add_argument('--config', dest='config', default='config_davis.json',
                    help='hyper-parameter of SiamMask in json format')
parser.add_argument('--base_path', default='../../data/tennis', help='datasets')
parser.add_argument('--cpu', action='store_true', help='cpu mode')
parser.add_argument('--track_mode', default='gt', choices = ['gt','sl'],help='track mode')
args = parser.parse_args()

if __name__ == '__main__':
    # Setup device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    torch.backends.cudnn.benchmark = True

    # Setup Model
    cfg = load_config(args)
    from custom import Custom
    siammask = Custom(anchors=cfg['anchors'])
    if args.resume:
        assert isfile(args.resume), 'Please download {} first.'.format(args.resume)
        siammask = load_pretrain(siammask, args.resume)

    siammask.eval().to(device)

    # Parse Image file
    img_files = sorted(glob.glob(join(args.base_path, '*.jp*')))
    ims = [cv2.imread(imf) for imf in img_files]

    # Select ROI
    cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN)
    # cv2.setWindowProperty("SiamMask", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    try:
        if args.track_mode == 'gt':
          gt, _ = os.path.split(args.base_path)
          gt += '//groundtruth.txt'
          f = open(gt, 'r')
          x1, y1, w1, h1 = f.readline().split(',')  
          x, y, w, h = int(x1),int(y1),int(w1),int(h1)
        else:
          init_rect = cv2.selectROI('SiamMask', ims[0], False, False)
          x, y, w, h = init_rect
          path+='//groundtruth.txt'
          x, y, w, h = str(x),str(y),str(w),str(h)
          f = open(path, 'w')
          f.write(x+','+y+','+w+','+h)
          f.close()

        cv2.waitKey(100)
    except:
        exit()


    ###
    gt, _ = os.path.split(args.base_path)
    path1, name = os.path.split(gt)
    file1 = open(path1 + '//track//' + name + '.txt', 'w')
    ###
    toc = 0
    for f, im in enumerate(ims):
        tic = cv2.getTickCount()
        if f == 0:  # init
            target_pos = np.array([x + w / 2, y + h / 2])
            target_sz = np.array([w, h])
            state = siamese_init(im, target_pos, target_sz, siammask, cfg['hp'], device=device)  # init tracker
            x, y, w, h = str(x), str(y), str(w), str(h)
            file1.write(x + ',' + y + ',' + w + ',' + h + '\n')
        elif f > 0:  # tracking
            state = siamese_track(state, im, mask_enable=True, refine_enable=True, device=device)  # track
            location = state['ploygon'].flatten()
            mask = state['mask'] > state['p'].seg_thr
            ######
            x, y = state['target_pos']
            w, h = state['target_sz']
            x = x - w / 2
            y = y - h / 2
            x, y, w, h = str(int(x)), str(int(y)), str(int(w)), str(int(h))
            if f != (len(ims) - 1):
                file1.write(x + ',' + y + ',' + w + ',' + h + '\n')
            else:
                file1.write(x + ',' + y + ',' + w + ',' + h)

                ######

            im[:, :, 2] = (mask > 0) * 255 + (mask == 0) * im[:, :, 2]
            #cv2.polylines(im, [np.int0(location).reshape((-1, 1, 2))], True, (0, 255, 0), 3)
            cv2.rectangle(im,(int(x),int(y)),(int(x)+int(w),int(y)+int(h)), color=(0,255,0))
            cv2.imshow('SiamMask', im)
            key = cv2.waitKey(1)

            if key > 0:
                break

        toc += cv2.getTickCount() - tic
    toc /= cv2.getTickFrequency()
    file1.close()
    fps = f / toc
    print('SiamMask Time: {:02.1f}s Speed: {:3.1f}fps (with visulization!)'.format(toc, fps))
